'''
Created on 2011-6-7

@author: elm
'''
import MySQLdb as mdb,string
from ddbms.Schema import *
from ddbms.SqlStatement import *
class Dataset(object):
    '''
    result set for a given query
    '''
    dataset = None
    types = {}
    
    def __init__(self,dataset=None,types=None):
        self.dataset = dataset
        self.types=types
        #self.checkDataType()
        
    def checkDataType(self):
        if self.types!=None and self.dataset!=None:            
            for d in self.dataset:
                for key in self.types:
                    if self.types[key]=="int":
                        if not isinstance(d[key],int):                            
                            d[key]=string.atoi(str(d[key]))
                            
        
    def getTypes(self):
        return self.types
    
    def data(self):
        return self.dataset
    
    def eval(self):
        return self
    
    def select(self,condition):
        ''
        results = []
        for row in self.data():
            if condition.toCNF().test(row):
                results += [row]
        return Dataset(results,self.types)
    
    def projection(self,fields):
        ''
        results = []
        types = {}
        for field in fields:
            fn = "%s.%s"%(field.table.table,field.column)
            if fn in self.types:
                types[fn] = self.types[fn]
        
        for row in self.data():
            result = {}
            for field in fields:
                fn = "%s.%s"%(field.table.table,field.column)
                if fn in row:
                    result[fn] = row[fn]
            results += [result]
        return Dataset(results,types)
    
    def join(self,key,otherDataset,otherKey):
        ''
        results = []
        types = self.types.copy()
        types.update(otherDataset.getTypes())
        d1 = self.data()
        d2 = otherDataset.data()
        if len(d1)==0 or len(d2)==0:
            return Dataset([],types)
        fnl = "%s.%s"%(key.table.table,key.column)
        fnr = "%s.%s"%(otherKey.table.table,otherKey.column)
        if types[fnr]=="int":
            for d in d2:
                if not isinstance(d[fnr],int):                    
                    d[fnr] = string.atoi(str(d[fnr]))                    
                            
        d2.sort(key=lambda x:x[fnr])
        def halffind(s,t,seq,key,keySelector):
            '''
            return first i that keySelector(seq[i])==key
            if i doen't exist, return None
            '''
            if s>t:
                return s
            else:
                m = (s+t)/2
                #if keySelector(seq[m])>=key:
                if seq[m][keySelector]>=key:
                    return halffind(s,m-1,seq,key,keySelector)
                else:
                    return halffind(m+1,t,seq,key,keySelector)
            
            
        for row1 in d1:
            k = row1[fnl]            
                            
            i = halffind(0,len(d2)-1,d2,k,fnr)
            if i>=0:
                l = len(d2)
                while i<l and d2[i][fnr] == k:
                    row2 = d2[i]                
                    result = row1.copy()
                    result.update(row2)
                    results += [result]
                    i+=1               
            
        return Dataset(results,types)
    
    def union(self,otherDataset):
        ''
        results = self.data() + otherDataset.data()
        return Dataset(results,self.types)
    def product(self,otherDataset):
        ''
        types = self.types.copy()
        types.update(otherDataset.getTypes())
        results = []
        d1 = self.data()
        d2 = otherDataset.data()
        for row1 in d1:
            result = row1.copy()
            for row2 in d2:
                result.update(row2)
                results += [result]
        return Dataset(results,types)
    def __str__(self):
        str = "%d tuples(%s):\n"%(len(self.dataset),self.types)
        if len(self.dataset)>10:
            for t in self.dataset[0:10]:
                str+="    %s\n"%t
        else:
            for t in self.dataset:
                str+="    %s\n"%t
        return str
    
class UnionedDeferredDataset(Dataset):
    'Dataset that unions several tables'
    
    sets = []
    
    mysql = None
    
    def __init__(self,sets,mysql):
        self.sets = []
        self.mysql = mysql
        for set in sets:
            if isinstance(set,UnionedDeferredDataset):
                self.sets += set.sets
            else:
                self.sets += [set]
    
    def getTypes(self):
        return self.sets[0].getTypes()
    
    def data(self):
        return self.eval().data()        
    
    def eval(self):
        defer = []
        evaled = []
        for ds in self.sets:
            if isinstance(ds,DeferredDataset):
                defer.append(ds.compile())
            else:
                evaled.append(ds)
        dds = None        
        if len(defer)>0:                    
            (host,user,pwd,db)=self.mysql
            conn = mdb.connect(host, user, pwd, db)
            cursor = conn.cursor(mdb.cursors.DictCursor)
            sql = " UNION ALL ".join("(%s)"%ds for ds in defer)
            cursor.execute(sql)
            dds = Dataset(list(cursor.fetchall()),self.getTypes())
        else:
            dds = Dataset([],self.getTypes())
        for d in evaled:
            dds = dds.union(d)
        return dds
    
    
    def select(self,condition):
        ds = []
        for d in self.sets:
            ds.append(d.select(condition))
        return UnionedDeferredDataset(ds)
    
    def projection(self,fields):
        ds = []
        for d in self.sets:
            ds.append(d.projection(fields))
        return UnionedDeferredDataset(ds,self.mysql)
        
    def join(self,key,other,otherKey):
        ds = []
        for d in self.sets:
            ds.append(d.join(key,other,otherKey))
        return UnionedDeferredDataset(ds,self.mysql)
        
    def union(self,other):
        if isinstance(other,UnionedDeferredDataset):
            ds = self.sets+other.sets
            return UnionedDeferredDataset(ds,self.mysql)
        else:
            ds = self.sets
            ds.append(other)
            return UnionedDeferredDataset(ds,self.mysql)
        
    def product(self,otherDataset):
        ds = []
        for d in self.sets:
            ds.append(d.product(otherDataset))
        return UnionedDeferredDataset(ds)
    
    def __str__(self):
        defer = []
        evaled = []
        for ds in self.sets:
            if isinstance(ds,DeferredDataset):
                defer.append(ds.compile())
            else:
                evaled.append(ds)        
        sql = " union ".join("(%s)"%ds for ds in defer)        
        if len(evaled)>0:
            sql += " And other %d datasets"%len(evaled)
        return sql
    
class DeferredDataset(Dataset):
    '''
    result set which is not fetched from mysqldb yet
    '''
    schema = None
    _select = None
    '[(Column,name)]'
    _from = None
    '[Segment]'
    _where = None
    'CNF'
    _sql = None
    mysql = None
    
    def __init__(self,schema,mysql):
        self.schema = schema
        self._select = []
        self._from = []
        self._where = []
        self.mysql = mysql
        
    def copy(self):
        dd = DeferredDataset(self.schema,self.mysql)
        dd._select = self._select
        dd._from = self._from
        dd._where = self._where
        dd._sql = self._sql
        return dd
    
    def data(self):
        (host,user,pwd,db)=self.mysql
        conn = mdb.connect(host, user, pwd, db)
        cursor = conn.cursor(mdb.cursors.DictCursor)
        cursor.execute(self.compile())
        d = list(cursor.fetchall())        
        return d
    
    def getTypes(self):
        types = {}
        for (col,name) in self._select:
            types[name]=col.type
        return types
    
    def eval(self):        
        return Dataset(self.data(),self.getTypes())
    
    def compile(self):
        if self._sql != None:
            return self._sql
        self._sql = "select %s from %s %s"%self.compilePart()
        return self._sql
    
    def col2Seg(self,col,segs):
            for seg in segs:
                if seg.columns.count(col.id)>0:
                    return seg
            return None
        
    def compilePart(self):        
        selectcol = set([])        
        fromtable = set([])
                
        for (col,name) in self._select:            
            fromseg = self.col2Seg(col,self._from)
            if fromseg != None:                
                fromtable.add(fromseg.name)
                #colexpr = "%s.%s as '%s'"%(fromseg.name,col.name,name)
                selectcol.add((fromseg.name,col.name,name))
        where = ""
        if isinstance(self._where,str):
            where = self._where
        elif(self._where!=None):
            cond = self._where.toCNF()
            cnf = []
            for dnf in cond.conditions:
                atoms = set([])
                for atom in dnf:
                    atom = atom.toSimple()
                    if isinstance(atom,TestValue):
                        segname = None
                        colname = None
                        if atom.fieldref.segmentRef:
                            fromtable.add(atom.fieldref.table.table)
                            segname = atom.fieldref.table.table
                            colname = atom.fieldref.column                            
                        else:
                            col = self.schema.findColumn(atom.fieldref)
                            seg = self.col2Seg(col,self._from)
                            fromtable.add(seg.name)                            
                            segname = seg.name
                            colname = col.name
                        val = None
                        if isinstance(atom.value,str):
                            val = "'%s'"%atom.value
                        else:
                            val = str(atom.value)
                        atomexpr = "(%s.%s %s %s)"%(segname,colname,atom.op,val)
                        atoms.add(atomexpr)
                    elif isinstance(atom,TestField):
                        lsegname = None
                        rsegname = None
                        lcolname = None
                        rcolname = None
                        if not atom.leftfieldref.segmentRef:
                            lcol = self.schema.findColumn(atom.leftfieldref)
                            lseg = self.col2Seg(lcol,self._from)
                            fromtable.add(lseg.name)
                            lsegname = lseg.name
                            lcolname = lcol.name
                        else:
                            fromtable.add(atom.leftfieldref.table.table)
                            lsegname = atom.leftfieldref.table.table
                            lcolname = atom.leftfieldref.column
                        if not atom.rightfieldref.segmentRef:
                            rcol = self.schema.findColumn(atom.rightfieldref)
                            rseg = self.col2Seg(rcol,self._from)
                            fromtable.add(rseg.name)
                            rsegname = rseg.name
                            rcolname = rcol.name
                        else:
                            fromtable.add(atom.rightfieldref.table.table)
                            rsegname = atom.rightfieldref.table.table
                            rcolname = atom.rightfieldref.column
                        
                        atomexpr = "(%s.%s %s %s.%s)"%(lsegname,lcolname,atom.op,rsegname,rcolname)
                        atoms.add(atomexpr)
                dnfexpr = None
                if len(atoms)>0:
                    dnfexpr = "(%s)"%(" OR ".join(atoms))
                else:
                    dnfexpr = "true"
                cnf.append(dnfexpr)
            where = "where (%s)"%(" AND ".join(cnf))
        selectcol = list(selectcol)
        selectcol.sort(key=lambda x:x[2])        
        cn = []
        for (a,b,c) in selectcol:
            cn.append("%s.%s as '%s'"%(a,b,c)) 
        return (",".join(cn),",".join(fromtable),where)
        #return self._sql
    
    def select(self,condition):
        ret = self.copy()
        if ret._where == None:
            ret._where = condition
        else:
            ret._where = LogicalAnd(ret._where,condition).toCNF()
        ret._sql = None         
        return ret
    
    def projection(self,fields):
        ret = self.copy()
        proj = []
        for fieldref in fields:
            proj.append(self.schema.findColumn(fieldref).id)
        cols = []
        for (col,name) in ret._select:
            if proj.count(col.id)>0:
                cols.append((col,name))
        ret._select = cols
        ret._sql = None  
        return ret
    
    def rewriteCond(self,cond,segs):
        if cond == None:
            return None
        cond = cond.toCNF()
        newcond = []
        for dnf in cond.conditions:
            dnfcond = []
            for atom in dnf:
                atom = atom.toSimple()
                if isinstance(atom,TestValue):
                    field = atom.fieldref                    
                    if not field.segmentRef:
                        col = self.schema.findColumn(field)
                        seg = self.col2Seg(col,segs)
                        field = FieldRef(TableRef(seg.name),field.column)
                        field.segmentRef = True                      
                    atom = TestValue(field,atom.op,atom.value)                        
                elif isinstance(atom,TestField):
                    lf = atom.leftfieldref
                    rf = atom.rightfieldref
                    if not lf.segmentRef:
                        lc = self.schema.findColumn(lf)
                        ls = self.col2Seg(lc,segs)
                        lf = FieldRef(TableRef(ls.name),lf.column)
                        lf.segmentRef = True
                    if not rf.segmentRef:
                        rc = self.schema.findColumn(rf)
                        rs = self.col2Seg(rc,segs)
                        rf = FieldRef(TableRef(rs.name),rf.column)
                        rf.segmentRef = True
                    atom = TestField(lf,atom.op,rf)
                else:
                    raise InvalidCondition
                dnfcond.append(atom)
            newcond.append(dnfcond)
        return CNF(newcond)
                    
    def join(self,key,other,otherKey):
        if isinstance(other,DeferredDataset):
            dd = DeferredDataset(self.schema,self.mysql)
            dd._select = self._select+other._select
            dd._from = self._from+other._from
            self._where = self.rewriteCond(self._where,self._from)
            other._where = other.rewriteCond(other._where,other._from)
            if self._where == None:
                if other._where == None:
                    if key != otherKey:
                        dd._where = TestField(key,Ops.EQ,otherKey).toCNF()
                    else:
                        leftseg = self.col2Seg(self.schema.findColumn(key), self._from)
                        rightseg = self.col2Seg(self.schema.findColumn(otherKey), other._from)
                        leftfield = FieldRef(TableRef(leftseg.name),key.column)
                        leftfield.segmentRef = True
                        rightfield = FieldRef(TableRef(rightseg.name),key.column)
                        rightfield.segmentRef = True
                        dd._where = TestField(leftfield,Ops.EQ,rightfield).toCNF()
                else:
                    if key!=otherKey:
                        dd._where = LogicalAnd(
                                        TestField(key,Ops.EQ,otherKey),
                                        other._where
                                    ).toCNF()
                    else:
                        leftseg = self.col2Seg(self.schema.findColumn(key), self._from)
                        rightseg = self.col2Seg(self.schema.findColumn(otherKey), other._from)
                        leftfield = FieldRef(TableRef(leftseg.name),key.column)
                        leftfield.segmentRef = True
                        rightfield = FieldRef(TableRef(rightseg.name),key.column)
                        rightfield.segmentRef = True
                        dd._where = LogicalAnd(
                                        TestField(leftfield,Ops.EQ,rightfield),
                                        other._where
                                    ).toCNF()
            else:
                if other._where == None:
                    if key!=otherKey:
                        dd._where = LogicalAnd(
                                        TestField(key,Ops.EQ,otherKey),
                                        self._where
                                    ).toCNF()
                    else:
                        leftseg = self.col2Seg(self.schema.findColumn(key), self._from)
                        rightseg = self.col2Seg(self.schema.findColumn(otherKey), other._from)
                        leftfield = FieldRef(TableRef(leftseg.name),key.column)
                        leftfield.segmentRef = True
                        rightfield = FieldRef(TableRef(rightseg.name),key.column)
                        rightfield.segmentRef = True
                        dd._where = LogicalAnd(
                                        TestField(leftfield,Ops.EQ,rightfield),
                                        self._where
                                    ).toCNF()
                else:
                    if key!=otherKey:
                        dd._where = LogicalAnd(
                                        self._where,       
                                        LogicalAnd(
                                            TestField(key,Ops.EQ,otherKey),
                                            other._where
                                    ))
                        dd._where = dd._where.toCNF()
                    else:
                        leftseg = self.col2Seg(self.schema.findColumn(key), self._from)
                        rightseg = self.col2Seg(self.schema.findColumn(otherKey), other._from)
                        leftfield = FieldRef(TableRef(leftseg.name),key.column)
                        leftfield.segmentRef = True
                        rightfield = FieldRef(TableRef(rightseg.name),key.column)
                        rightfield.segmentRef = True
                        dd._where = LogicalAnd(
                                        self._where,       
                                        LogicalAnd(
                                            TestField(leftfield,Ops.EQ,rightfield),
                                            other._where
                                    )).toCNF()
            return dd
        else:
            return other.join(otherKey,self,key)
        
    def union(self,other):        
        if isinstance(other,DeferredDataset):
            ds = [self ,other]
            return UnionedDeferredDataset(ds,self.mysql)
        else:
            return other.union(self)
        
    def product(self,other):        
        if isinstance(other,DeferredDataset):
            self._where = self.rewriteCond(self._where, self._from)
            other._where = self.rewriteCond(other._where, other._from)
            dd = DeferredDataset(self.schema)
            dd._select = self._select+other._select
            dd._from = self._from+other._from
            if self._where == None:
                dd._where = other._where
            else:
                if other._where == None:
                    dd._where = self._where
                else:
                    dd._where = LogicalAnd(self._where,other._where).toCNF()
            return dd
        else:
            return other.product(self)
    
    def __str__(self):
        return "[DeferredDataSet: %s]"%self.compile()
